from langchain.llms import Ollama
from langchain.vectorstores import Chroma
# from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from pathlib import Path
import requests
import json
from config import OLLAMA_CONFIG
from models.document import Knowledge
from chromadb import PersistentClient
from config import CHROMA_CONFIG
from chromadb.config import Settings

class ChatService:
    def __init__(self):
        self.base_url = OLLAMA_CONFIG['base_url']
        
    def get_available_models(self):
        """获取 Ollama 可用的模型列表"""
        try:
            response = requests.get(f"{self.base_url}/api/tags")
            if response.status_code == 200:
                models = response.json().get('models', [])
                return [{'id': model['name'], 'name': model['name']} for model in models]
            return []
        except Exception as e:
            print(f"Error fetching models: {e}")
            return []

    def get_embeddings(self, model_name):
        """
        获取embedding模型实例
        Args:
            model_name (str): Ollama中的embedding模型名称
        Returns:
            OllamaEmbeddings: embedding模型实例
        """
        from langchain.embeddings import OllamaEmbeddings
        return OllamaEmbeddings(
            base_url=self.base_url,
            model=model_name
        )
    
    def get_llm(self, model_name="llama2"):
        """获取语言模型实例"""
        return Ollama(base_url=self.base_url, model=model_name)
    
    def create_chain(self, collection_name, model_name="llama2", embedding_model=None, top_k=3):
        """
        创建对话链
        Args:
            collection_name (str): 知识库集合名称
            model_name (str): 对话模型名称
            embedding_model (str): embedding模型名称
            top_k (int): 检索的文档数量
        """
        vector_db = Chroma(
            client=PersistentClient(path=str(Path(CHROMA_CONFIG['path']).resolve()), settings=Settings(anonymized_telemetry=False)),
            collection_name=collection_name,
            embedding_function=self.get_embeddings(embedding_model)
        )
        
        # 使用传入的 top_k 参数
        retriever = vector_db.as_retriever(search_kwargs={"k": top_k})
        
        memory = ConversationBufferMemory(
            memory_key="chat_history",
            output_key="answer",  # 添加这行，指定输出键
            return_messages=True
        )
        
        chain = ConversationalRetrievalChain.from_llm(
            llm=self.get_llm(model_name),
            retriever=retriever,
            memory=memory,
            return_source_documents=True,
            chain_type="stuff"  # 添加这行，使用默认的chain类型
        )
        
        return chain

    def get_streaming_chain(self, collection_name, model_name="llama2", embedding_model=None, top_k=3):
        """创建支持流式响应的对话链"""
        retriever = self.get_retriever(
            collection_name=collection_name,
            embedding_model=embedding_model,
            top_k=top_k
        )
        relevant_docs = retriever.get_relevant_documents
        return relevant_docs, self.get_llm(model_name)

    def get_retriever(self, collection_name, embedding_model, top_k=3):
        """
        获取检索器
        Args:
            collection_name (str): 集合名称
            embedding_model (str): embedding模型名称
            top_k (int): 检索的文档数量
        Returns:
            Retriever: 文档检索器
        """
        vector_db = Chroma(
            client=PersistentClient(path=str(Path(CHROMA_CONFIG['path']).resolve()), settings=Settings(anonymized_telemetry=False)),
            collection_name=collection_name,
            embedding_function=self.get_embeddings(embedding_model)
        )
        return vector_db.as_retriever(search_kwargs={"k": top_k})

    def stream_chat(self, question: str, knowledge: Knowledge, model_name: str, top_k: int = 3):
        """流式聊天"""
        try:
            # 使用新的 get_retriever 方法
            retriever = self.get_retriever(
                collection_name=knowledge.collection_name,
                embedding_model=knowledge.embedding_model,
                top_k=top_k
            )
            docs = retriever.get_relevant_documents(question)
            
            if not docs:
                yield "我无法从知识库中找到相关信息。"
                return
                
            # 构建上下文
            context = "\n".join([doc.page_content for doc in docs])
            
            # 构建 prompt
            prompt = f"""基于以下内容回答问题。如果无法从提供的内容中找到答案，请说'我无法从提供的内容中找到答案'。

文档内容：
{context}

问题：{question}

回答："""
            
            # 调用 Ollama 流式 API
            response = requests.post(
                f"{self.base_url}/api/generate",
                json={"model": model_name, "prompt": prompt, "stream": True},
                stream=True
            )
            
            for line in response.iter_lines():
                if not line:
                    # 显式发送换行符
                    yield "\n"
                    continue
                
                try:
                    data = json.loads(line.decode())
                    if "response" in data:
                        text = data["response"]
                        # 确保换行符被正确传递
                        if text.strip():  # 只处理非空内容
                            yield text
                except Exception as e:
                    print(f"Error processing stream: {e}")
                    continue
                    
        except Exception as e:
            yield f"发生错误: {str(e)}"

def generate_streaming_response(chat_service, question, knowledge, model_name, top_k):
    try:
        retriever, llm = chat_service.get_streaming_chain(
            collection_name=knowledge.collection_name,
            model_name=model_name,
            embedding_model=knowledge.embedding_model,
            top_k=top_k
        )
        
        # 获取相关文档
        docs = retriever(question)
        context = "\n".join([doc.page_content for doc in docs])
        
        # 构建 prompt
        prompt = f"""基于以下内容回答问题。如果无法从提供的内容中找到答案，请说'我无法从提供的内容中找到答案'。

文档内容：
{context}

问题：{question}

回答："""
        
        # 使用 Ollama 的流式响应
        response = requests.post(
            f"{OLLAMA_CONFIG['base_url']}/api/generate/stream",
            json={"model": model_name, "prompt": prompt, "stream": True},
            stream=True
        )
        
        for line in response.iter_lines():
            if not line:
                continue
                
            try:
                # 解码二进制数据为字符串
                line_str = line.decode('utf-8')
                # 移除可能的"data: "前缀
                if line_str.startswith("data: "):
                    line_str = line_str[6:]
                    
                data = json.loads(line_str)
                
                if "response" in data:
                    yield f"data: {data['response']}\n\n"
                if data.get("done", False):
                    yield "data: [DONE]\n\n"
                    break
            except json.JSONDecodeError as e:
                print(f"JSON decode error: {e}, line: {line_str}")
                continue
            except Exception as e:
                print(f"Error processing line: {e}")
                continue
                    
    except Exception as e:
        print(f"Streaming error: {e}")
        yield f"data: Error occurred: {str(e)}\n\n"
        yield "data: [DONE]\n\n"
